package ai.koog.agents.core.agent

import ai.koog.agents.core.agent.AIAgent.Companion.State
import ai.koog.agents.core.agent.AIAgent.Companion.State.NotStarted
import ai.koog.agents.core.agent.context.AIAgentContext
import ai.koog.agents.core.agent.context.element.AgentRunInfoContextElement
import ai.koog.agents.core.agent.entity.AIAgentStrategy
import ai.koog.agents.core.feature.AIAgentFeature
import ai.koog.agents.core.feature.pipeline.AIAgentPipeline
import io.github.oshai.kotlinlogging.KLogger
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.withContext
import kotlin.reflect.KClass
import kotlin.uuid.ExperimentalUuidApi
import kotlin.uuid.Uuid

/**
 * Abstract base class representing a single-use AI agent with state.
 *
 * This AI agent is designed to execute a specific long-running strategy only once and provides an API to monitor and manage its state.
 *
 * It maintains internal states including its running status, whether it was started, its result (if available), and
 * the root context associated with its execution. The class enforces safe state transitions and provides
 * thread-safe operations via a mutex.
 *
 * @param Input the type of the input accepted by the agent.
 * @param Output the type of the output produced by the agent.
 * @param TContext the type of the context used during the agent's execution, extending [AIAgentContext].
 * @property logger the logger used for logging execution details and errors.
 * @param id the unique identifier for the agent. Random UUID will be generated if set to null.
 */
@OptIn(ExperimentalUuidApi::class)
public abstract class StatefulSingleUseAIAgent<Input, Output, TContext : AIAgentContext>(
    protected val logger: KLogger,
    id: String? = null,
) : AIAgent<Input, Output> {
    /**
     * A mutex used to synchronize access to the state of the agent. Ensures that only one coroutine
     * can modify or read the shared state of the agent at a time, preventing data races and ensuring
     * thread-safe operations.
     */
    private val agentStateMutex: Mutex = Mutex()

    private var state: State<Output> = NotStarted()

    final override suspend fun getState(): State<Output> = agentStateMutex.withLock { state.copy() }

    final override val id: String by lazy { id ?: Uuid.random().toString() }

    /**
     * The execution strategy defining how the agent processes input and produces output.
     */
    public abstract val strategy: AIAgentStrategy<Input, Output, TContext>

    /**
     * Represents the pipeline used by the AI agent for processing tasks or data.
     *
     * This abstract property defines the structure or sequence of operations
     * within the AI agent's pipeline. It serves as the core mechanism for
     * executing workflows, handling inputs, and generating outputs in the
     * AI agent's functionality.
     */
    protected abstract val pipeline: AIAgentPipeline

    /**
     * Executes the agent's main functionality, coordinating with various components
     * such as pipelines and strategies. Ensures the agent is run in a thread-safe
     * manner using locking mechanisms.
     *
     * @param agentInput The input required to execute the agent's strategy.
     *                   This includes any data necessary for processing.
     * @return The output generated by the agent's execution, produced as a
     *         result of applying the strategy to the provided input.
     * @throws IllegalStateException if the agent was already started.
     * @throws Throwable if any exception occurs during the execution process.
     */
    final override suspend fun run(agentInput: Input): Output {
        agentStateMutex.withLock {
            if (state !is NotStarted) {
                throw IllegalStateException(
                    "Agent was already started. Please use AIAgentService.createAgentAndRun(agentInput) to run an agent multiple times."
                )
            }
            state = State.Starting()
        }

        val runId = Uuid.random().toString()
        val context = prepareContext(agentInput, runId)

        return withContext(
            AgentRunInfoContextElement(
                agentId = id,
                runId = runId,
                agentConfig = agentConfig,
                strategyName = strategy.name
            )
        ) {
            pipeline.withPreparedPipeline {
                agentStateMutex.withLock {
                    state = State.Running(context)
                }

                logger.debug {
                    formatLog(
                        agentId = id,
                        runId = runId,
                        message = "Starting agent execution"
                    )
                }

                pipeline.onAgentStarting<Input, Output>(
                    runId = runId,
                    agent = this@StatefulSingleUseAIAgent,
                    context = context
                )

                val result = try {
                    strategy.execute(context = context, input = agentInput)
                } catch (e: Throwable) {
                    logger.error(e) { "Execution exception reported by server!" }
                    pipeline.onAgentExecutionFailed(
                        agentId = id,
                        runId = runId,
                        throwable = e
                    )
                    agentStateMutex.withLock { state = State.Failed(e) }
                    throw e
                }

                logger.debug {
                    formatLog(
                        agentId = id,
                        runId = runId,
                        message = "Finished agent execution"
                    )
                }
                pipeline.onAgentCompleted(
                    agentId = id,
                    runId = runId,
                    result = result
                )

                agentStateMutex.withLock {
                    state = if (result != null) {
                        State.Finished(result)
                    } else {
                        State.Failed(Exception("result is null"))
                    }
                }

                result ?: error("result is null")
            }
        }
    }

    /**
     * Closes the AI Agent and performs necessary cleanup operations.
     *
     * This method is a suspending function that ensures that the AI Agent's resources are released
     * when it is no longer required. It notifies the pipeline of the agent's closure and ensures
     * that any associated features or stream providers are properly closed.
     *
     * Overrides the `close` method to implement agent-specific shutdown logic.
     */
    final override suspend fun close() {
        // TODO: Delete Closeable interface from [AIAgent] declaration.
    }

    /**
     * Prepares and initializes the agent context required to handle the given input and run ID.
     *
     * @param agentInput the input provided to the agent for processing.
     * @param runId a unique identifier representing the current execution or operation run.
     * @return the initialized context specific to the agent setup for the provided input and run ID.
     */
    public abstract suspend fun prepareContext(agentInput: Input, runId: String): TContext

    /**
     * Retrieves a feature from the agent's pipeline associated with this agent using the specified key.
     *
     * @param TFeature A feature implementation type.
     * @param feature A feature to fetch.
     * @param featureClass The [KClass] of the feature to be retrieved.
     * @return The feature associated with the provided key, or null if no matching feature is found.
     * @throws IllegalArgumentException if the specified [featureClass] does not correspond to a registered feature.
     */
    public fun <TFeature : Any> feature(
        featureClass: KClass<TFeature>,
        feature: AIAgentFeature<*, TFeature>
    ): TFeature? = pipeline.feature(featureClass, feature)

    /**
     * Formats a log message with the specified agent ID, run ID, and message content.
     *
     * @param agentId The unique identifier of the agent generating the log.
     * @param runId The unique identifier of the specific run or task associated with the log.
     * @param message The content of the log message to be formatted.
     * @return A formatted log string containing the agent ID, run ID, and the provided message.
     */
    protected fun formatLog(agentId: String, runId: String, message: String): String =
        "[agent id: $agentId, run id: $runId] $message"

    /**
     * Executes the provided block within a prepared pipeline context.
     * Ensures that the necessary feature resources are initialized before the block is invoked
     * and properly cleaned up after finishing the block.
     *
     * @return The result of the block's execution.
     */
    private suspend fun <T> AIAgentPipeline.withPreparedPipeline(block: suspend () -> T): T =
        try {
            prepareAllFeatures()
            block.invoke()
        } finally {
            onAgentClosing(agentId = id)
            closeAllFeaturesMessageProcessors()
        }
}

/**
 * Retrieves a feature from the [StatefulSingleUseAIAgent.pipeline] associated with this agent using the specified key.
 *
 * @param feature A feature to fetch.
 * @return The feature associated with the provided key, or null if no matching feature is found.
 * @throws IllegalArgumentException if the specified [feature] does not correspond to a registered feature.
 */
public inline fun <reified TFeature : Any> StatefulSingleUseAIAgent<*, *, *>.feature(
    feature: AIAgentFeature<*, TFeature>
): TFeature? = feature(TFeature::class, feature)

/**
 * Retrieves a feature from the [StatefulSingleUseAIAgent.pipeline] associated with this agent using the specified key
 * or throws an exception if it is not available.
 *
 * @param feature A feature to fetch.
 * @return The feature associated with the provided key
 * @throws IllegalStateException if the [TFeature] feature does not correspond to a registered feature.
 * @throws NoSuchElementException if the feature is not found.
 */
public inline fun <reified TFeature : Any> StatefulSingleUseAIAgent<*, *, *>.featureOrThrow(
    feature: AIAgentFeature<*, TFeature>
): TFeature = feature(feature) ?: throw NoSuchElementException("Feature ${feature.key} is not found.")
